#ifndef SPRITE_RENDERER_H
#define SPRITE_RENDERER_H

#include "sprite.h"
#include "ecs.h"

#include <csugl.h>

#include <glad/glad.h>

#include <array>

struct SpriteRenderer
{
    explicit SpriteRenderer(int32_t sprite_capacity, int32_t texture_capacity)
    {
        _shader = SGT_Ref_Str_AM(csugl::Shader)->get("sprite_shader");
        _batch_data = csugl::MakeRef<BatchData>(sprite_capacity * 4, sprite_capacity * 4);
        _texture_buffer = csugl::MakeRef<TextureBuffer>(texture_capacity);
    }

    inline int sprite_capacity() {return batch_capacity() / 4;}

    inline int batch_capacity() { return static_cast<int>(_batch_data->vertex_cap); }

    inline int batch_nums() { return _batch_data->batch_count; }

    void begin(csugl::Camera* camera) {
        _shader->use();
        _shader->set_mat4("_V", camera->get_view_mat());
        _shader->set_mat4("_P", camera->get_projection_mat());
        _batch_data->start();
    }

    void submit(const Sprite& sprite, const csugl::Transform& trans) {
        int texture_ID = this->_texture_buffer->add(sprite.texture);
        if (texture_ID < 0)
        {
            this->flush();
            this->_texture_buffer->restart();
            texture_ID = this->_texture_buffer->add(sprite.texture);
        }

        auto ts_f = static_cast<glm::vec2>(glm::ivec2{sprite.texture->width, sprite.texture->height});
        glm::mat4 mat({
            0,      ts_f.y, 0.f, 1.f,
            0,      0,      0.f, 1.f,
            ts_f.x, 0,      0.f, 1.f,
            ts_f.x, ts_f.y, 0.f, 1.f
        });
        auto mat_trans = trans.get_trans_mat() * mat;
        this->_batch_data->fill({
            //pos                               // texCoords                           // texture ID
            {{mat_trans[0].x, mat_trans[0].y}, {sprite.uv_start.x, sprite.uv_end.y},   texture_ID},
            {{mat_trans[1].x, mat_trans[1].y}, {sprite.uv_start.x, sprite.uv_start.y}, texture_ID},
            {{mat_trans[2].x, mat_trans[2].y}, {sprite.uv_end.x, sprite.uv_start.y},   texture_ID},
            {{mat_trans[3].x, mat_trans[3].y}, {sprite.uv_end.x, sprite.uv_end.y},     texture_ID}
        }, {
            {0,1,2},
            {0,3,2}
        });
    }

    void flush()
    {
        _texture_buffer->bind(_shader);
        _batch_data->flush();
    }

    void end() {
        this->flush();
        _batch_data->restart();
        _texture_buffer->restart();
    }

private:
    
    struct SpriteVertex
    {
        glm::vec2 position;
        glm::vec2 texCoords;
        // glm::vec4 color;
        int tex_id;
    };
    
    struct BatchData
    {
        explicit BatchData(uint64_t vertex_capacity, uint64_t triangle_capacity)
            : vertex_cap(vertex_capacity),
              triangle_cap(triangle_capacity),
              vertices(new SpriteVertex[vertex_capacity]), v_end(vertices),
              indices(new glm::u32vec3[triangle_capacity]), i_end(indices)
        {
            vao = csugl::VertexArray::Creat();
            vbo = csugl::VertexBuffer::Create<SpriteVertex>(
                vertex_cap,
                nullptr,
                csugl::BufferUsage::DYNAMIC_DRAW);
            vao->addVertexAttribution(vbo, {{
                csugl::BufferElementType::Float2,
                csugl::BufferElementType::Float2,
                csugl::BufferElementType::Int,
            }});
            ibo = csugl::IndexBuffer::Create<glm::u32vec3>(
                triangle_cap,
                nullptr,
                csugl::BufferUsage::DYNAMIC_DRAW);
        }

        inline void start() 
        {
            batch_count = 0;
            v_end = vertices;
            i_end = indices;
        }

        inline void restart() 
        {
            v_end = vertices;
            i_end = indices;
        }

        void fill(std::initializer_list<SpriteVertex> vert,
                  std::initializer_list<glm::u32vec3> indi)
        {
            if (v_end - vertices > static_cast<int32_t>(vertex_cap - vert.size()))
            {
                flush();
            }
            for (auto &v : vert)
            {
                *v_end = v;
                v_end++;
            }
            glm::u32vec3 offset{static_cast<uint32_t>(v_end - vertices - vert.size())};
            for (auto &i : indi)
            {
                *i_end = i + offset;
                i_end++;
            }
        }

        void flush()
        {
            auto i_count = 3 * (i_end - indices);
            if (i_count == 0)
                return;
            vao->bind();
            vbo->setData(0, v_end - vertices, vertices);
            ibo->setData(0, i_end - indices, indices);
            glDrawElements(GL_TRIANGLES, static_cast<GLsizei>(i_count), GL_UNSIGNED_INT,
                           const_cast<const void *>(static_cast<void *>(0)));
            batch_count++;
            restart();
        }

        int32_t batch_count = 0;
        uint64_t vertex_cap, triangle_cap;
        csugl::Ref<csugl::VertexArray> vao;
        csugl::Ref<csugl::VertexBuffer> vbo;
        csugl::Ref<csugl::IndexBuffer> ibo;

    private:
        friend SpriteRenderer;

        SpriteVertex *vertices, *v_end;
        glm::u32vec3 *indices, *i_end;
    };

    struct TextureBuffer
    {
        TextureBuffer(int capacity)
            : capacity(capacity), textures(new csugl::Ref<csugl::Texture2D>[capacity])
        {
        }

        void restart()
        {
            for (int i = 0; i < capacity; i++)
            {
                textures[i].reset();
            }
        }

        int add(csugl::Ref<csugl::Texture2D> newTex)
        {
            for (int i = 0; i < capacity; i++)
            {
                if (textures[i] == newTex) // find
                    return i;
                if (nullptr == textures[i]) // new
                {
                    textures[i] = newTex;
                    return i;
                }
            }
            return -1;
        }

        void bind(csugl::Ref<csugl::Shader> shader)
        {
            for (int i = 0; i < capacity; i++)
            {
                if (textures[i])
                {
                    csugl::TextureUnit t_u;
                    t_u.active(i);
                    textures[i]->bind();
                    std::stringstream prop_str;
                    prop_str << "_texture" << i;
                    shader->set_sampler2D(prop_str.str(), i);
                }
                else
                {
                    return;
                }
            }
        }

        int capacity;
        csugl::Ref<csugl::Texture2D> *textures;
    };

    csugl::Ref<csugl::Shader> _shader;
    csugl::Ref<BatchData> _batch_data;
    csugl::Ref<TextureBuffer> _texture_buffer;

};

struct SpriteRendererComponent : public Component
{
    SpriteRendererComponent(const Sprite &sprite, int sot_layer) : sprite(sprite) {}
    SpriteRendererComponent(Sprite &&sprite, int sot_layer) : sprite(sprite) {}
    Sprite sprite;
    int sort_layer = 0;
};

struct SpriteRendererSystem : public System
{
    SpriteRendererSystem(csugl::Ref<Scene> cur_scene) : System(cur_scene), sprd(1000, 8)  {}

    void start() override
    {
        auto win = &csugl::singleton<csugl::Application>::getInstance()->GetWindow();
        win->addEventCallback(
            [this](auto &ev)
            {
                csugl::EventDispatcher d(ev);
                d.Dispatch<csugl::WindowResizeEvent>(
                    [&, this](auto &ev)
                    {
                        _cur_scene->ComponentEach<CameraComponent>(
                            [&, this](const auto &ca_c)
                            {
                                auto ca = dynamic_cast<csugl::OrthographicCamera *>(ca_c.camera.get());
                                CHECK(ca) << "cast to OrthographicCamera failed";
                                ca->aspect = static_cast<float>(ev.height) / ev.width;
                            });
                        return false;
                    });
            });
    }

    void update() override
    {
        _cur_scene->SortAs<SpriteRendererComponent>(
            [](const SpriteRendererComponent &sp_c_a, const SpriteRendererComponent &sp_c_b)
            {
                return sp_c_a.sort_layer > sp_c_b.sort_layer;
            });
        _cur_scene->ComponentEach<CameraComponent, TransformComponent>([&](CameraComponent& camera_c, const TransformComponent& trans_c){
            camera_c.camera->transform = trans_c.transform;
            sprd.begin(camera_c.camera.get());
        });
        _cur_scene->ComponentEach<SpriteRendererComponent, TransformComponent, ActiveComponent>(
            [&](const auto &sprite_c, const auto &trans_c, const auto &active_c)
            {
                if(active_c.active)
                    sprd.submit(sprite_c.sprite, trans_c.transform);
            });
        sprd.end();
    }

    int batch_nums() { return sprd.batch_nums(); }

private:
    SpriteRenderer sprd;
};

#endif